from keras.models import Sequential
from keras.layers import Conv2D, MaxPooling2D
from keras.layers import Flatten, Dense, Dropout
from keras.preprocessing.image import ImageDataGenerator

def classification (batch_size = 20, epochs = 20, img_width = 150,
                    img_height = 150, num_train_samples = 2000,
                    num_validation_samples = 800):

    model = Sequential()
    model.add(Conv2D(32, (3, 3), activation='relu',
              input_shape=(img_width, img_height, 3)))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Conv2D(48, (3, 3), activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Conv2D(64, (3, 3), activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.4))      # Dropout
    model.add(Flatten()) 
    model.add(Dense(128, activation='relu'))
    model.add(Dense(1, activation='sigmoid'))

    model.compile(loss='binary_crossentropy',
              optimizer='rmsprop',
              metrics=['accuracy'])
    
    print(model.summary())
    
    # augmentation processes for the train set
    train_image_generator = ImageDataGenerator (  
        rescale = 1./255.,
        rotation_range = 25,
        width_shift_range = 0.2,
        height_shift_range = 0.2,
        zoom_range = 0.3,
        horizontal_flip = True,
        fill_mode = 'nearest'
    )
    
    train_image_generator.fit(X_train)

    # test images are not augmented
    validation_image_generator = ImageDataGenerator (  
        rescale = 1./255.
    )
    
    batch_size = 32
    train_generator = train_image_generator.flow(X_train, y_train, 
                                                 batch_size = batch_size)
    validation_generator  = validation_image_generator.flow(X_test,
                                         y_test, batch_size = batch_size)
       
    history = model.fit_generator(
        train_generator,
        # 8000 training images in groups of batch_size
        steps_per_epoch = num_train_samples // batch_size,    
        epochs = 20,
        validation_data = validation_generator,
        # 1200 testing images in groups of batch_size 
        validation_steps = num_validation_samples // batch_size,   
        verbose = 0)
    
    return history
